#!/usr/bin/env python3
import sys
import re
from datetime import datetime

def main():
    if len(sys.argv) != 2:
        print("用法: {} <日志文件>".format(sys.argv[0]))
        sys.exit(1)
    log_file = sys.argv[1]

    # 从文件名中提取任务开始时间，要求文件名格式为 train_YYYYMMDD_HHMMSS.log
    pattern_filename = r"train_(\d{8})_(\d{6})\.log"
    match_filename = re.search(pattern_filename, log_file)
    if not match_filename:
        print("日志文件名格式错误，应为 train_YYYYMMDD_HHMMSS.log")
        sys.exit(1)
    date_str = match_filename.group(1)
    time_str = match_filename.group(2)
    try:
        start_time = datetime.strptime(date_str + time_str, "%Y%m%d%H%M%S")
    except Exception as e:
        print("解析开始时间失败:", e)
        sys.exit(1)

    # 定义正则：匹配包含 "datetime:" 的行，提取[]中的阶段描述和时间戳（格式为YYYY-MM-DD HH:MM:SS）
    pattern_line = re.compile(r"\[(.*?)\]\s*datetime:\s*(\d{4}-\d{2}-\d{2}\s+\d{2}:\d{2}:\d{2})")
    events = []
    with open(log_file, "r") as f:
        for line in f:
            match_line = pattern_line.search(line)
            if match_line:
                stage = match_line.group(1).strip()
                timestamp_str = match_line.group(2).strip()
                try:
                    event_time = datetime.strptime(timestamp_str, "%Y-%m-%d %H:%M:%S")
                except Exception as e:
                    print("解析时间失败:", e)
                    continue
                elapsed = (event_time - start_time).total_seconds()
                events.append((stage, event_time, elapsed))

    if not events:
        print("日志文件中没有找到包含 'datetime:' 的时间点")
        sys.exit(1)

    # 输出统计结果
    print("任务开始时间（来自文件名）:", start_time)
    print("\n各阶段时间统计：")
    for stage, event_time, elapsed in events:
        print(f"阶段: {stage}\n  时间: {event_time}\n  耗时: {elapsed:.0f} 秒\n")
    
    # 根据最后一个时间点计算任务总耗时
    total_elapsed = events[-1][2]
    print("任务总耗时（秒）:", int(total_elapsed))

if __name__ == '__main__':
    main()
